use crate::stream::{CommandStream, DataStream};
use std::collections::{HashMap, HashSet};
use std::net::TcpStream;
use std::thread;
use std::sync::{Arc, Mutex};
use std::process::exit;

pub struct Client {
    host: String,
    port: u16,
    cmd_stream: CommandStream,
    proxy_map: Arc<Mutex<HashMap<u16, u16>>>
}

impl Client {
    pub fn new(host: &str, port: u16, proxy_ports: &HashSet<(u16, u16)>) -> Client {
        let tcp = match TcpStream::connect((host, port)) {
            Ok(s) => s,
            Err(e) => {
                eprintln!("ERROR: {}", e);
                exit(-1);
            }
        };
        let mut cmd_stream = CommandStream::new(tcp);
        let mut set_client = HashSet::<u16>::new();
        let mut set_server = HashSet::<u16>::new();
        let mut proxy_map = HashMap::<u16, u16>::new();
        for p in proxy_ports {
            if set_client.contains(&p.0) { continue; }
            if set_server.contains(&p.1) { continue; }
            set_client.insert(p.0);
            set_server.insert(p.1);
            proxy_map.insert(p.0, p.1);
            cmd_stream.write(format!("listen {}", p.1))
                .expect("Connection error");
        }
        cmd_stream.write(String::from("end"))
            .expect("Connection error");
        println!("Successfully connect to: {}:{}", host, port);
        Client {
            host: host.to_string(),
            port,
            cmd_stream,
            proxy_map: Arc::new(Mutex::new(proxy_map))
        }
    }

    pub fn run(&mut self) {
        loop {
            let cmd = match self.cmd_stream.read() {
                Ok(s) => s,
                Err(_) => {
                    eprintln!("Server closed connection");
                    std::process::exit(0);
                }
            };
            self.exec_cmd(cmd);
        }
    }

    pub fn exec_cmd(&mut self, cmd: String) {
        let mut args = cmd.split(' ');
        let arg0 = match args.next() {
            Some(a) => a,
            None => { return; }
        };
        match arg0 {
            "request" => {
                let arg1 = match args.next() {
                    Some(a) => a,
                    None => { return; }
                };
                let arg2 = match args.next() {
                    Some(a) => a,
                    None => { return; }
                };
                self.request(arg1, arg2);
            }
            _ => { eprintln!("ERROR: '{}' command not found", arg0); }
        }
    }

    fn request(&mut self, arg1: &str, arg2: &str) {
        let port = match arg1.parse::<u16>() {
            Ok(p) => p,
            Err(_) => { return; }
        };
        let index = match arg2.parse::<u32>() {
            Ok(p) => p,
            Err(_) => { return; }
        };
        let map = self.proxy_map.clone();
        let host = self.host.clone();
        let host_port = self.port;
        thread::spawn(move || {
            let server_port = port;
            let mut port = 0;

            let locked_map = match map.lock() {
                Ok(l) => l,
                Err(_) => { return; }
            };

            let mut find = false;
            for en in locked_map.iter() {
                if *en.1 == server_port {
                    find = true;
                    port = *en.0;
                    break;
                }
            }
            if !find { return; }

            let tcp_stream = match TcpStream::connect(("127.0.0.1", port)) {
                Ok(s) => s,
                Err(_) => { return; }
            };
            let stream = match TcpStream::connect((host.as_str(), host_port)) {
                Ok(s) => s,
                Err(_) => { return; }
            };
            let mut data_stream = DataStream::new(stream);
            match data_stream.write_head(format!("{} {}", server_port, index)) {
                Ok(_) => {},
                Err(_) => { return; }
            };
            data_stream.pipe(tcp_stream);
        });
    }
}
